Skip to content

[PyTorch/Jax] Fix attention mask definition, and sliding window for decoder - #818

Merged
ptrendx merged 29 commits into
NVIDIA:mainfrom
cyanguwa:fix_attn_masks
May 17, 2024
Merged

[PyTorch/Jax] Fix attention mask definition, and sliding window for decoder#818
ptrendx merged 29 commits into
NVIDIA:mainfrom
cyanguwa:fix_attn_masks

Conversation

@cyanguwa

@cyanguwa cyanguwa commented Apr 26, 2024 •
edited by ptrendx
Loading

Copy link
Copy Markdown
Collaborator

Description

Fixes #614.
Fixes #629.

Moving forward, we'd like to define attention mask consistently in PyTorch, Jax and Paddle as True being masking out the corresponding position and False being allowing that position to participate in attention. A typical causal mask would look like:

0, 1, 1, 1
0, 0, 1, 1
0, 0, 0, 1
0, 0, 0, 0

Type of change

  • Documentation change (change only to the documentation, either a fix or a new content)
  • Bug fix (non-breaking change which fixes an issue)
  • New feature (non-breaking change which adds functionality)
  • Breaking change (fix or feature that would cause existing functionality to not work as expected)

Changes

  • Unify the implementation and unit tests regarding attention mask in TE PyTorch, Jax and Paddle.
  • Differentiate sliding widow initialization for decoder and encoder
  • Save aux_ctx_tensors and similar tensors using save_for_backward call instead of to ctx
  • Skip Jax tests for fused attention backward when bias shape is not 1HSS

Checklist:

  • I have read and followed the contributing guidelines
  • The functionality is complete
  • I have commented my code, particularly in hard-to-understand areas
  • I have made corresponding changes to the documentation
  • My changes generate no new warnings
  • I have added tests that prove my fix is effective or that my feature works
  • New and existing unit tests pass locally with my changes

cyanguwa added 3 commits April 26, 2024 22:37
@cyanguwa
fix inconsistency for attn mask; now True means participating in attn
fae6572 
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
@cyanguwa
fix sliding window window_size for decoder+padding combination
dc6869d 
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
@cyanguwa
Merge branch 'NVIDIA:main' into fix_attn_masks
cca5f49
@timmoon10

timmoon10 commented Apr 29, 2024

Copy link
Copy Markdown
Member

For reference, most frameworks treat attention mask Trues as inclusion:

Implementation True in attention mask implies inclusion
torch.nn.MultiheadAttention No
torch.nn.functional.scaled_dot_product_attention Yes
flax.linen.MultiHeadAttention Yes
paddle.nn.MultiHeadAttention Yes
keras.layers.Attention Yes
tf.keras.layers.Attention Yes
cuDNN fused attention Yes
apex.contrib.multihead_attn.SelfMultiheadAttn No
flash_attn.fused_softmax.ScaledMaskedSoftmax No

@cyanguwa

cyanguwa commented Apr 30, 2024

Copy link
Copy Markdown
Collaborator Author

@timmoon10 thanks for the table 😺 I was doing something like that myself (OCD? not if someone else is also doing it!).

After looking through some of the other frameworks' implementations (PyTorch/Jax/Paddle), and our own implementation and unit tests, I'm leaning towards True = masking out; False = inclusion in attention actually.

I don't think Keras is a popular framework now and that we should follow its trend. Also, PyTorch's scaled_dot_product_attention is a beta feature but its MultiHeadAttention is stable. cuDNN's attention doesn't really use a mask in my opinion, not the True/False kind anyway. For causal, it's using a is_causal=True/False flag, and for padding, a cu_seqlens tensor (which we converted the mask to in TE).

Our fused softmax implementation uses True = masking out; False = inclusion in attention, which also makes all the softmax calls in TE PyTorch, Jax and Paddle all subscribe to that definition. They are not the fastest attention paths, so maybe we don't have to care about them as much. But, I think the original term when it came out, for causal, is upper triangle masking, which kind of made it upper rather than lower, and that 1 means masking things out. This is probably also why apex and Tri Dao attention's softmax both adopted that definition.

Any thoughts? 😃 thanks!

cyanguwa added 3 commits May 1, 2024 22:52
@cyanguwa
revert paddle changes regarding mask
81c4605 
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
@cyanguwa
Merge branch 'main' into fix_attn_masks
35de3be
@cyanguwa
revert softmax to 1-mask;0-keep
4ca194c 
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
@cyanguwa cyanguwa changed the title [PyTorch] Fix attention mask definition and sliding window for decoder [PyTorch/Jax] Fix attention mask definition and sliding window for decoder May 1, 2024
cyanguwa added 9 commits May 2, 2024 00:39
@cyanguwa
enforce 1-mask out; 0-keep rule for jax masks
3a4b377 
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
@cyanguwa
fix jax lint
dfa0ece 
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
@cyanguwa
revert pytorch mask changes; some kept in tests
aa6eaca 
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
@cyanguwa
revert to jax fused attn on main
44062a9 
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
@cyanguwa
inverse mask logic for get_cu_seqlens/_and_indices in PyTorch impleme…
1c8a073 
…ntation and mask generation in unit tests

Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
@cyanguwa
temporarily disable update_weight_scale_inv
19f271c 
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
@cyanguwa
enforce window_size for decoder
2a489bf 
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
@cyanguwa
add docstring for mask definition 1-mask out;0-keep
87d02b6 
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
@cyanguwa
Merge branch 'main' into fix_attn_masks
ea06868
@cyanguwa

cyanguwa commented May 2, 2024

Copy link
Copy Markdown
Collaborator Author

/te-ci

@cyanguwa
cyanguwa requested review from timmoon10 and zlsh80826 May 2, 2024 22:54
@zlsh80826

zlsh80826 commented May 3, 2024

Copy link
Copy Markdown
Collaborator

Hi @cyanguwa, thanks for making things clear. TE-JAX awares mask_bit = True to mask out, and it is documented, so I think JAX's unit tests don't need to be changed in this PR.

zlsh80826
zlsh80826 reviewed May 3, 2024
View reviewed changes
Comment thread tests/jax/test_fused_attn.py
@ptrendx ptrendx added the 1.7.0 label May 3, 2024
@Marks101

Marks101 commented May 7, 2024

Copy link
Copy Markdown
Contributor

@cyanguwa Thank you so much for looking into this and clarifying the mask definition ❤️

cyanguwa added 3 commits May 13, 2024 17:12
@cyanguwa
Merge branch 'main' into fix_attn_masks
0489d74 
Signed-off-by: cyanguwa <8636796+cyanguwa@users.noreply.github.com>
@cyanguwa
add aux_ctx_tensors to save_for_backward
d759aee 
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
@cyanguwa
Merge branch 'main' into fix_attn_masks
6b42b93
@cyanguwa cyanguwa changed the title [PyTorch/Jax] Fix attention mask definition and sliding window for decoder [PyTorch/Jax] Fix attention mask definition, and sliding window for decoder May 14, 2024
@cyanguwa

cyanguwa commented May 14, 2024

Copy link
Copy Markdown
Collaborator Author

/te-ci pytorch

@cyanguwa
tweak make_decoder_mask and make_mask in jax tests
2e84099 
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
@cyanguwa

cyanguwa commented May 14, 2024

Copy link
Copy Markdown
Collaborator Author

/te-ci jax

@cyanguwa
cyanguwa requested a review from zlsh80826 May 14, 2024 22:23
zlsh80826
zlsh80826 approved these changes May 15, 2024
View reviewed changes

@zlsh80826 zlsh80826 left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM!

cyanguwa added 3 commits May 15, 2024 17:30
@cyanguwa
skip dBias for shapes other than 1HSS; otherwise dq/dk/dv NaNs
49a38f0 
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
@cyanguwa
expand attn_biases from list to variables in save_for_backward
5bd7a1a 
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
@cyanguwa
fix use of variable before assignment in jax dact_lu
bf9a851 
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
@cyanguwa

cyanguwa commented May 15, 2024

Copy link
Copy Markdown
Collaborator Author

/te-ci pytorch

@cyanguwa

cyanguwa commented May 15, 2024

Copy link
Copy Markdown
Collaborator Author

/te-ci jax

cyanguwa added 6 commits May 15, 2024 21:24
@cyanguwa
remove window size definition for decoder
84aa282 
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
@cyanguwa
Merge branch 'main' into fix_attn_masks
3911898
@cyanguwa
add change notes in README for padding mask in PyTorch
37f25c3 
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
@cyanguwa
tweak padding mask notes in README
87cea82 
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
@cyanguwa
Merge branch 'main' into fix_attn_masks
5e65f15
@cyanguwa
expand list to tensors for save_for_backwards
f07b0bf 
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
@cyanguwa

cyanguwa commented May 16, 2024

Copy link
Copy Markdown
Collaborator Author

/te-ci pytorch

ptrendx
ptrendx approved these changes May 17, 2024
View reviewed changes

@ptrendx ptrendx left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM

@ptrendx
ptrendx merged commit 67bc399 into NVIDIA:main May 17, 2024
This was referenced Jul 28, 2026
Closed
Merged
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@ptrendx ptrendx ptrendx approved these changes
@timmoon10 timmoon10 Awaiting requested review from timmoon10
+1 more reviewer
@zlsh80826 zlsh80826 zlsh80826 approved these changes
Reviewers whose approvals may not affect merge requirements

Assignees

No one assigned

Labels

1.7.0

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

[PyTorch] FlashAttention: causal masking enforced in cross attention due to sliding window attention [PyTorch] Definition of attention mask

5 participants

@cyanguwa @timmoon10 @zlsh80826 @Marks101 @ptrendx